What

Under shard_mode: explicit, MaxText pays a step-time penalty that traces to a single unfoldable transpose on the LM head's weight gradient. This adds lm_head_weight_grad_in_kernel_order (default off), which removes it, and drops the previously-proposed lm_head_kernel_transposed in favour of it.

Why

JAX's lower_with_sharding_in_types annotates every sharding-in-types primitive's output with a custom_call_target="Sharding" op — 65 of them under auto on llama2 versus 1,276 under explicit. One lands between the LM head weight-gradient dot and the transpose that follows it, and XLA's algebraic simplifier will not fold transpose(dot(A,B)) → dot(B,A) across it. Census on the dp4 module: auto has 17 adjacent (foldable) transpose-of-dot pairs, explicit has 0.

The transpose survives to codegen and costs twice:

1. A lost reduce-scatter. The gradient is left sharded on its minor-most dimension, below the ~512-byte per-shard extent tpu-all-reduce-scatter-fusion requires, so the weight-gradient all-reduce is never rewritten into a reduce-scatter — 2× the wire bytes for that tensor.
2. Six relayout copies per step. The resulting layout propagates through the whole Adam update, so the donated optimizer state must be converted in and out — 196.6 MB/step on llama2, up to 1.57 GB/step at emb 4096.

The barrier itself cannot be removed: it is JAX's, not MaxText's. Gating MaxText's out_sharding pin produces a 0-line HLO diff.

How

DenseGeneral gains weight_grad_in_kernel_order, backed by a custom_vjp whose backward pass contracts dk straight into the kernel's stored axis order — so there is no transpose to fold in the first place. The stored kernel, its sharding and its initialization are untouched, which makes this checkpoint-compatible; the gradient differs from the default rule only by floating-point reassociation.

It self-disables under shard_mode: auto (where XLA folds the transpose on its own) and under quantization (where it cannot own the dot). logits_via_embedding is rejected rather than silently ignored. The untied LM head opts in via lm_head_weight_grad_in_kernel_order on both the linen and nnx decoder paths.

Measured

v5p, 4 chips, xprof train_step_ns median on /device:TPU:0, 56 runs. All arms run back-to-back per config so within-row deltas carry no session drift.

Please read the two halves of that table differently — this is the part that is easy to over-claim.

The five emb 512 rows go from a +4.59% median penalty to −0.57%, beating auto outright on three of five. But emb 512 / fsdp 4 puts 256 bytes per shard on the minor-most dimension, below XLA's 512-byte AR→RS gate, so mechanism (1) is active only there. No production job has that geometry — those rows are an artifact of the shrunk benchmark proxies.

The two bold rows are the real ship criterion, and the bar there is no regression, not a win: prod (emb 2048, 16 layers, seq 4k) is +0.03% against a 0.03% rep spread, w2048 is −0.57%. There was almost nothing left to recover at that geometry, and recovering it costs nothing.

Acceptance

On llama2/shrink/fsdp4 the explicit dump gains %all-reduce-scatter.2, byte-for-byte identical to auto's %all-reduce-scatter.5 — bf16[512,32000] → bf16[128,32000], layout {1,0:T(8,128)(2,1)}, dim-0 dynamic-slice. f32[emb/n, 32000] relayout copies in after_optimizations drop from 6 to 0, matching auto.

Numerical equivalence

20 training steps against the auto baseline, same seed, max absolute loss deviation:

The flag sits at explicit's own noise level — on mixtral slightly below it — which is what reassociation alone predicts.

Changed since the first revision of this PR

lm_head_kernel_transposed has been removed. That flag stored the kernel as [vocab, embed] to reach the same layout, and it did work. It is dropped because the custom_vjp dominates it on every axis measured:

• Slower on four of seven configs by 1.85–4.02%, because it recovers only the relayout copies and not the reduce-scatter, and within noise on the two where it already won.
• No checkpoint-conversion path — it changes the on-disk shape of params-decoder-logits_dense-kernel, so it cannot be flipped on an existing run.
• Changes initialization. At the flipped shape jax.random draws a different matrix from the same seed, so it perturbs the loss ~35× more than explicit's own noise (llama2 9.42e-03 vs 2.52e-04) — reinitialization, not rounding.

Shipping a strictly-dominated flag that also carries a silent checkpoint hazard seemed worse than shipping one flag that works, so the transposed-kernel code, its config plumbing, its integration/tunix and vLLM converter changes, and its tests are all gone. Happy to restore it if you would rather have both.

A second negative result is documented rather than shipped. Eliding redundant reshard barriers (no-op reshards where the aval already carries the requested sharding) looked like the obvious companion fix. It was implemented and measured: it removes ~4% of barriers in a real module and the optimized executable is op-identical (llama2 2889→2889, prod 2942→2942, deepseek 9642→9642, zero differing ops). XLA already discards no-op Sharding custom-calls. Recorded in the doc so the next person does not spend the week.

Testing

• tests/unit/linears_test.py — 7 new tests: kernel bit-identity, forward equality, gradient equality across three contraction shapes, the LM-head 2-D shape in fp32 and bf16, composition with nnx.remat, the auto-mode no-op, and the quantization yield. Whole file passes (21 passed, 1 TPU-only skipped).
• tests/unit/pyconfig_test.py, tests/unit/configs_value_test.py — pass.
• End-to-end train on v5p, 6 steps, flag off vs on: loss identical to 3 decimals at every step.
• pyink clean, pylint clean (the two remaining findings in types.py pre-exist on main).

Notes

• Default stays false. The flag is inert under auto and checkpoint-safe under explicit, so flipping the default is a follow-up decision, not a prerequisite.
• Full analysis, including the mechanism, the causal controls, and both negative results, is in docs/guides/optimization/shard_mode_performance.md.
• The largest remaining lever measured but not in this PR: extending the same custom_vjp to the in-loop DenseGenerals, gated on a leading FSDP axis. Estimated +360.7 µs at prod geometry, but it is untested under nnx.scan + remat and is roughly a week of work.

Follow-up (2026-09-03)

Four commits pushed on top of the above. Two of them close the last items this PR left open — the default, and the in-loop DenseGenerals — and two are separate wins found while measuring. Both "Notes" bullets above are superseded: the default is now on, and the in-loop lever is in this PR.

The one table

Eight onboarded models at emb 2048 / mlp 8192 / 16 layers (deepseek 8, gemma3 18 — it scans in groups of 6), seq 1024, pdbs 1, fsdp 4 on v5p-8. One build, 54 runs, 3 reps per arm, medians. The explicit column has nothing in the config but shard_mode: explicit — it is what a user gets.

Median −0.330%, faster on five of eight and inside the rep spread on a sixth. With the two per-model dense-flag settings applied, the worst regression anywhere is mixtral's +0.041%, which is inside mixtral's own 0.07% rep spread. Before this work the worst cases were qwen3 +3.25%, gemma3 +2.02% and deepseek +0.92%.

lm_head_weight_grad_in_kernel_order is now on by default where it can act

Re-measured on every model, at L12/L14/L16, and at vocab 4096 and 32000, the flag has never lost to auto — there is no measured configuration where off is the better choice, so off was the wrong default.

Mixtral is the one model the flag does not help; its −1.393% is an MoE-scan scheduling windfall unrelated to the LM head (it oscillates +0.120% / −0.151% / −0.038% across L12/L14/L16, and shrinking the vocab 8× leaves the spread unchanged), and the flag costs it nothing — it lands on auto's schedule rather than losing to it.

The flag is therefore tri-state: None resolves to on for an untied head under shard_mode: explicit, off everywhere else. An explicit true on a tied head is now an error rather than silently ignored.

dense_weight_grad_in_kernel_order — the in-loop lever, shipped default-off

The same custom_vjp wired to the three MlpBlock, four attention and nine MLA/Indexer projections. It turned out not to be a one-signed missed optimization. Under explicit the Sharding custom-calls decouple the gradient's layout from the stored parameter's, so XLA picks the dot's preferred layout; the flag declines that freedom and reproduces auto's layout op for op. On the deep-scan models both modes relayout anyway and explicit's source is 24% cheaper to convert (20.2 vs 26.4 µs/copy), so the flag gives back a real win:

The discriminator is the scanned gradient stack's depth — the two winners are 3 and 1 deep, the six losers all 16 — but two points below the line against six above is not enough to fit a default on, so it ships off, with the two models named in base.yml and in section 4.8 of the doc. Stored kernels and their initialization are untouched, so there is no checkpoint hazard, and it is inert under auto and on quantized layers.

Two wins that are not about sharding at all

YarnRotaryEmbedding built its whole frequency table inside every layer of every step to read max_target_length rows out of it — f32[163840, 32] on deepseek, 0.6% of it used. Computing the rows straight from position is −1.4% of step time, equally in both modes, which is the check that it is not a sharding effect. It survived a full mode-comparison sweep precisely because both modes paid it: A/B diffing two modes cannot find work both modes do. The detector that found it — flag any op whose output element count is ≥ 8× its largest input's — was then run over all eight models and found nothing else above ~0.06%.

attend_on_embedding materialized table.T. Under explicit that is a distinct typed value XLA cannot fold into the dot's dimension numbers, so the table shard was cast to bf16 twice and the tied head's gradient came out flipped. Expressed as einsum dimension numbers instead.

Also in the doc, not in the code

• Section 6.3, explicit's largest and most stable win: explicit deletes 480 of 495 bf16[24,8192] collective-permute-start ops — the rotated-reduce-scatter fixup — on every untied model.
• Thirteen refuted candidates, each implemented or measured rather than reasoned about, so the next person does not spend the time.
• A systematic hunt for any remaining explicit-only cost on llama2/mistral/qwen3 above the gate found nothing above ~0.10% of step.

Testing (follow-up commits)

• tests/unit/linears_test.py, tests/unit/embeddings_test.py, tests/unit/pyconfig_test.py — 72 passed, 1 TPU-only skipped, 26 subtests. New coverage: the custom_vjp on >2-D kernels via _permuted_sharding, freqs_cis_at bit-identical to indexing the full table at the bottom/middle/top of a 163,840-row table, attend_on_embedding forward and gradient equality across three query ranks, and the tri-state resolution including the tied-head rejection.
• tests/unit/configs_test.py, tests/unit/configs_value_test.py — 113 passed, 110 subtests.
• pyink, pylint, codespell, mdformat, yamllint clean via pre-commit (the two types.py pylint findings pre-exist on main).

Second follow-up (2026-09-03): the depth sweep that invalidated the tables above

Three commits: 5b9725702, 6a54684a7, 169aac772.

Everything above this line was measured at 4, 8 or 16 decoder layers. All three are scan depths where an unrelated XLA layout pathology inflates both sharding modes by ~25%, and a 25% inflation is more than enough scheduling noise to swamp — and repeatedly flip the sign of — a 0.4% sharding effect. The tables are kept and marked rather than deleted, because the corrections reverse conclusions that looked solid.

The pathology (doc section 4.9)

Walking llama2 from 2 to 24 layers at emb 2048 / mlp 8192, ns-per-layer is flat at ~3.3 M except at stack lengths 2, 4, 8, 16 and 24, where it jumps 25–34%:

At those lengths XLA assigns layout {2,0,1} to the scanned parameter stack f32[L,512,8192] instead of {2,1,0}, so every dynamic-update-slice writes a degenerate T(1,128)-tiled slice — three op families costing 8.16 ms of a 36.5 ms step at L8, byte-identical in both sharding modes. param_scan_axis: 0 removes it (L8 −23.6%, L16 −24.9%) and costs 0.18–0.28% at healthy depths; it survives at production geometry (+14.1% at L8, +11.5% at L16). This is unrelated to shard_mode, is larger than everything else in the doc, and is not changed by this PR — it is documented as section 8 item 03 for someone to take on separately.

What that does to the two flags

Re-measured where the pathology does not fire, dense_weight_grad_in_kernel_order has one sign everywhere and the two flags are additive. llama2 / qwen3 at 18 layers, four ways:

Across 15 (model, depth) pairs the LM-head flag alone leaves a median +0.463% penalty against auto (worst +1.038%, gemma2); with both flags the median is +0.005% and the full range is −0.183% … +0.134%. The "two models only" recommendation in the section above was fitted entirely to contaminated points, so dense_weight_grad_in_kernel_order now defaults on under shard_mode: explicit, same tri-state shape as the LM-head flag.

The replacement table (doc section 5.6)

Eight models, d16 geometry, 3 reps, nothing in the config but shard_mode: explicit — both flags resolve themselves. Every depth chosen so the pathology is not firing:

Median +0.000%; worst case +0.133% against a 0.16% rep spread on the same row. Six of eight rows sit inside their own rep spread in both directions; the two that do not are gemma-2b and gemma3-4b, and explicit wins both.

The three models this round set out to fix, before → after:

gemma3 is where the dense flag does real work: 1.18 pp, and it changes the sign. mixtral and deepseek were already inside noise and still are; what the flag buys there is not a win but the absence of a tail.

mixtral and deepseek are at 14 and 12 layers because at 18 both exceed HBM at this geometry (deepseek asks 43.56 G of temporaries against 31.24 G available). Both are healthy stack lengths, which is the property that matters.

custom_vjp → jax.sharding.auto_axes

Both flags are now implemented by tracing the forward dot inside a jax.sharding.auto_axes region rather than by hand-writing the reordered gradient. The region drops the sharding barrier for the length of that one dot, so Shardy re-propagates the operands' shardings exactly as shard_mode: auto would and XLA performs the fold itself:

def _dot_general_in_auto_axes(inputs, kernel, dimension_numbers, precision, out_sharding):
  dot = functools.partial(lax.dot_general, dimension_numbers=dimension_numbers, precision=precision)
  if out_sharding is None:
    out_sharding = jax.eval_shape(dot, inputs, kernel).sharding
    if out_sharding is None:  # no mesh -> no barrier to remove
      return dot(inputs, kernel)
  return jax.sharding.auto_axes(dot, out_sharding=out_sharding)(inputs, kernel)

114 fewer lines, gradients now bit-identical to the default rule rather than agreeing to rounding, no operand-permutation bookkeeping for >2-D kernels, and it works at all 17 call sites including MLA. Step times agree with the custom_vjp to within 0.15 pp on all five models measured both ways.

The unit test for the mechanism changed accordingly: auto_axes leaves the transpose in the jaxpr (XLA folds it later), so the test now asserts the barrier is gone — under explicit the weight-gradient transpose's operand is defined by stablehlo.dot_general with the flag on and by sdy.sharding_constraint with it off.

Disclosed, not fixed

• mixtral prefers explicit with neither flag, by ~1.4% at L12/L14/L16. It reproduces, survives an 8× vocab shrink, and has the opposite sign on the other MoE model in the sweep. The shipped default gives that up to land on parity; lm_head_weight_grad_in_kernel_order: false is called out in the doc for mixtral users who can measure it.
• param_scan_axis: 0 is not made a default here. It is worth ~25% at the affected depths but costs 0.18–0.28% elsewhere and has only been characterized on one model family and one chip generation.

Testing

• tests/unit/linears_test.py, tests/unit/pyconfig_test.py — 60 passed, 1 TPU-only skipped, 20 subtests.
• tests/unit/sharding_test.py, sharding_nnx_test.py, sharding_compare_test.py, sharding_desc_test.py, moe_test.py, attention_test.py — 190 passed, 126 skipped.
• pyink, pylint (10.00/10 on both changed Python files), codespell, mdformat, yamllint clean via pre-commit.